package com.stress.testing.plugin.apache.rocketmq.hook;

import com.stress.testing.common.TraceContext;
import com.stress.testing.common.TraceContextHolder;
import org.apache.rocketmq.client.hook.SendMessageContext;
import org.apache.rocketmq.client.hook.SendMessageHook;
import org.apache.rocketmq.common.message.Message;

public class PluginSendMessageHook implements SendMessageHook {
    @Override
    public String hookName() {
        return PluginSendMessageHook.class.getSimpleName();
    }

    @Override
    public void sendMessageBefore(SendMessageContext context) {
        if (context == null) {
            return;
        }

        Message msg = context.getMessage();
        if (msg == null) {
            return;
        }

        TraceContext traceContext = TraceContextHolder.getTraceContext();
        if (traceContext == null) {
            return;
        }

        // 添加链路信息
        for (String name : traceContext.propertyNames()) {
            msg.putUserProperty(name, traceContext.getProperty(name));
        }
    }

    @Override
    public void sendMessageAfter(SendMessageContext context) {

    }
}
